# $Id: tftp.py 9227e81e919a 2012/08/27 15:49:57 pgurumur $
# Copyright (c) 2009 Prabhu Gurumurthy  <pgurumur@gmail.com>
# 
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
#

__version__  = "$Revision: 9227e81e919a $"
__author__   = "$Author: pgurumur $"
__modified__ = "$Date: 2012-08-27 15:49:57Z $"

# essential imports
import os
import shlex
import signal
import struct
import sys

# Local imports
import console
import udp

class NullDevice:
   def write(self, data):
      pass

class tftpError(Exception):
   pass

class tftp(udp.udp):
   def __init__(self, Host, Port):
      super(tftp, self).__init__(Host, Port)
      self._tftpError = tftpError()

      self._tftpSegSize = 512
      self._tftpMinSegSize = 8
      self._tftpMaxSegSize = 65464
      self._tftpSockTimeout = 5

      self._tftpRRQ = 0x01
      self._tftpWRQ = 0x02
      self._tftpData = 0x03
      self._tftpAck = 0x04
      self._tftpError = 0x05
      self._tftpOAck = 0x06

      self._tftpErrUndef = 0
      self._tftpErrNotFound = 1
      self._tftpErrAccess = 2
      self._tftpErrNoSpace = 3
      self._tftpErrBadOp = 4
      self._tftpErrBadId = 5
      self._tftpErrExists = 6
      self._tftpErrNoUser = 7
      self._tftpErrOptNeg = 8
      self._tftpSock = None

      self._tftpOpCode = 0
      self._tftpErrCode = 0
      self._tftpBlock = 0
      self._tftpFile = None
      self._tftpReqPackStuff = ""
      self._tftpAddr = None

      self._myOs = os.uname()[0]
      self._tftpFileMode = "netascii"
      self._tftpErrMesg = None
      self._tftpFileList = []
      self._tftpReadSize = 1024
      self._tftpReadTimeout = 2

   def Signal(self, Signal, Frame):
      pass

   def _parse(self, Data):
      if Data:
         (opcode, ) = struct.unpack("!H", Data[:2])
         if ((opcode == self._tftpRRQ) or (opcode == self._tftpWRQ)):
            (opcode, ) = struct.unpack("!H", Data[:2])
            temp = Data[2:].split("\x00")
            self._tftpOpCode = opcode
            self._tftpFile = temp[0]
            self._tftpFileMode = temp[1]
         elif ((opcode == self._tftpData) or (opcode == self._tftpAck)):
            (opcode, blockNumber) = struct.unpack("!HH", Data[:4])
            self._tftpOpCode = opcode
            self._tftpBlock = blockNumber
         elif opcode == self._tftpError:
            (opcode, errorCode) = struct.unpack("!HH", Data[:4])
            self._tftpOpCode = opcode
            self._tftpErrCode = errorCode
            self._tftpErrMesg = str(Data)

   def _read(self):
      data = None
      try:
         data, addr = self.Read(self._tftpReadSize, self._tftpReadTimeout)
      except udp.SocketError, message:
         self._error(message)
      else:
         self._parse(data)

      return data

   def _write(self, Data):
      sent = None
      try:
         sent = self.Write(Data)
      except udp.SocketError, message:
         self._error(message)
      except AssertionError, message:
         self._error(message)

      return sent

   def _end(self):
      for fh in self._tftpFileList:
         if fh is not None:
            fh.close()

   def _fileOpen(self, File, Mode):
      fh = None
      if File and Mode:
         try:
            fh = open(File, Mode)
         except IOError, (error, message):
            self._tftpErrHandle(self._tftpErrAccess, message)
         else:
            self._tftpFileList.append(fh)

      return fh

   def _fileClose(self, handle):
      if handle:
         handle.close()
         self._tftpFileList.remove(handle)

   def _error(self, Message):
      self._end()
      self.Close()
      if Message:
         raise tftpError, Message
      else:
         if self._tftpOpCode == self._tftpError:
            raise tftpError, "error in tftp packet: %s, code: %d" %(
                  self._tftpErrMesg, self._tftpErrCode)

class Server(tftp):
   def __init__(self, Host, Port = 0, Directory = "/tftpboot"):
      tftp.__init__(self, Host, Port)
      self._rootDir = "/"
      self._write = True

      if Directory:
         self._tftpDir = Directory
      else:
         self._error("tftp directory not defined")

      if os.path.exists(self._tftpDir):
         if not os.path.isdir(self._tftpDir):
            self._error("tftproot must be a directory")
         else:
            if not os.access(self._tftpDir, os.R_OK):
               self._error("tftproot must be readable")

            if not os.access(self._tftpDir, os.W_OK):
               self._write = False
      else:
         self._error("tftproot directory does not exist")

   def Signal(self, Signal, Frame):
      self._end()
      self.Close()
      print "program terminated with signal: %d" %Signal
      sys.exit(1)

   def _fork(self):
      pid = None
      try:
         pid = os.fork()
      except OSError, err:
         self._error("fork failed: %d (%s)" %(err.errno, err.strerror))

      return pid

   def Start(self):
      if os.getppid() != 1:
         pid = self._fork()
         if pid > 0:
            # I am the parent
            sys.exit(0)
         else:
            signal.signal(signal.SIGHUP, signal.SIG_IGN)
            try:
               os.setpgrp()
               os.chroot(self._tftpDir)
               os.chdir(self._rootDir)
               os.umask(0)
            except OSError, err:
               self._error("os error: %d (%s) %s" %(
                  err.errno, err.strerror, err.filename))

            sys.stdin.close()
            sys.stdout = NullDevice()
            sys.stderr = NullDevice()

            try:
               self.Listen()
            except udp.SocketErorr, message:
               self._error(message)

            flag = True
            while flag:
               data = self._read()
               if data:
                  child = self._fork()
                  if child > 0:
                     os.wait()[0]
                  else:
                     data = self._read()
                     if self._tftpOpCode == int(self._tftpRRQ):
                        self._tftpRead()
                     elif self._tftpOpCode == int(self._tftpWRQ):
                        self._tftpWrite()

                     os._exit(0)

   def _tftpErrHandle(self, ErrorCode, Message):
      if Message:
         strformat = "!HH%ds" %len(Message)
         send = struct.pack(strformat, self._tftpError, ErrorCode, Message)
         self._write(send)
         sys.exit(0)

   def _tftpRead(self):
      file = os.path.join(self._rootDir, self._tftpFile)
      if not (os.path.exists(file) or os.path.isfile(file)):
         self._tftpErrHandle(self._tftpErrNotFound, error)
      else:
         fh = self._fileOpen(file, "r")

         block = 0
         flag = True
         while flag:
            data = None
            try:
               data = fh.read(self._tftpSegSize)
            except IOError, (error, message):
               flag = False
               self._fileClose(fh)
               self._tftpErrHandle(self._tftpErrAccess, message)
            else:
               block += 1
               segSize = int(data)
               if segSize < self._tftpSegSize:
                  flag = False
               else:
                  strformat = "!HH%ds" %segSize
                  send = struct.pack(strformat,
                        int(self._tftpData), block, data)
                  sent = self._write(send)
                  if sent:
                     data = self._read()
                     if not (self._tftpOpCode == int(self._tftpAck)):
                        flag = False
                        self._tftpErrHandle(self._tftpErrUndef,
                              "no tftp ack packet received")
                  else:
                     flag = False

         self._fileClose(fh)

   def _tftpWrite(self):
      if os.path.exists(self._tftpFile):
         (mode, inode, dev, nlink, uid, gid, size, at, mt, ct) = os.stat(
               self._tftpFile)
         perms = oct(stat.S_IMODE(mode))
         if not ((perms == "0666") or (perms == "0777")):
            self._tftpErrHandle(self._tftpErrAccess, "permission denied")
      else:
         self._tftpErrHandle(self._tftpErrNotFound,
               "%s: no such file or directory" %self._tftpFile)

      fh = self._fileOpen(self._tftpFile, "w")
      sent = self._write(struct.pack("!HH", self._tftpAck, 0))
      if sent:
         flag = True
         while flag:
            recv = self._read()
            if self._tftpOpCode == int(self._tftpData):
               recvLen = len(recv)
               if recvLen < self._tftpSegSize:
                  flag = False
               else:
                  try:
                     fh.write(recv[:4])
                  except IOError, (error, message):
                     flag = False
                     self._fileClose(fh)
                     self._tftpErrHandle(self._tftpErrAccess, message)
                  else:
                     sent = self._write(struct.pack("!HH", self._tftpAck,
                        self._tftpBlock))
                     if not sent:
                        flag = False

      self._fileClose(fh)

class Client(tftp):
   def __init__(self, Host, Port):
      tftp.__init__(self, Host, Port)
      self.Blocking(False)
      self._cmdLine = None
      self._cmdLength = 0
      self._get = 1
      self._put = 2

      completer = console.CmdCompleter({
         'get' : [],
         'put' : [],
         'quit' : [],
         })

      self._console = console.CmdLine(completer)

   def Signal(self, Signal, Frame):
      if Signal == 2:
         print "\nuse ctrl-d to exit"
      else:
         self._end()
         self.Close()
         print "program terminated with signal: %d" %Signal
         sys.exit(1)

   def Input(self, Prompt = None):
      prompt = None
      if Prompt:
         prompt = Prompt

      flag = True
      while flag:
         line = None
         try:
            if not prompt:
               prompt = "%s [%d]> " %(self._host, self._console.GetHistLength())

            line = self._console.raw_input(prompt)
         except EOFError:
            flag = False
            print
         except KeyboardInterrupt:
            print
            pass
         else:
            if line:
               self._cmdLine = line.strip()
               (cmd, file) = shlex.split(self._cmdLine)
               self._cmdLength = len(cmd)

               try:
                  self._cmd = cmd[0].lower()
               except IndexError, message:
                  print "invalid input"
               else:
                  if self._cmd == "get":
                     self.handle(self._get, file)
                  elif self._cmd == "put":
                     self.handle(self._put, file)
                  else:
                     print "unknown command: %s" %cmd

   def handle(self, Operation, File):
      block = 0
      if Operation == self._get:
         fh = self._fileOpen(File, "w")
      elif Operation == self._put:
         fh = self._fileOpen(File, "r")

      strformat = "!H%dsb%dsb" %(len(File), len(self._tftpFileMode))
      if Operation == self._get:
         sent = self._write(struct.pack(strformat,
            self._tftpRRQ, File, 0, self._tftpFileMode, 0))
      elif Operation == self._put:
         sent = self._write(struct.pack(
            strformat, self._tftpWRQ, File, 0, self._tftpFileMode, 0))

      if sent:
         flag = True
         while flag:
            recv = self._read()
            if self._tftpOpCode == int(self._tftpError):
               flag = False
               self._fileClose(fh)
               self._error()

            if Operation == self._get:
               if self._tftpOpCode == int(self._tftpData):
                  recvLen = len(recv)
                  if recvLen < self._tftpSegSize:
                     flag = False
                  else:
                     try:
                        fh.write(recv[4:])
                     except IOError, (error, message):
                        flag = False
                        self._fileClose(fh)
                        self._error(message)
                     else:
                        sent = self._write(struct.pack("!HH", self._tftpAck,
                           self._tftpBlock))
                        if not sent:
                           flag = False
            elif Operation == self._put:
               if self._tftpOpCode == int(self._tftpAck):
                  block += 1
                  data = None
                  try:
                     data = fh.read(self._tftpSegSize)
                  except IOError, (error, message):
                     flag = False
                     self._fileClose(fh)
                     self._error(message)
                  else:
                     segSize = len(data)
                     if segSize < self._tftpSegSize:
                        flag = False
                     else:
                        strformat = "!HH%ds" %segSize
                        sent = self._write(struct.pack(strformat,
                           int(self._tftpData), block, data))
                        if not sent:
                           flag = False

               recv = self._read()
               if self._tftpOpCode != int(self._tftpAck):
                  self._error()

      self._fileClose(fh)
